{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ef0e97e0-a18b-4a4d-b895-a212045084ea",
   "metadata": {},
   "source": [
    "# AutoGen Tool calling capabilities with Unity Catalog\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "**API Key**\n",
    "To run this tutorial, you will need an OpenAI API key. \n",
    "\n",
    "Once you have acquired your key, set it to the environment variable `OPENAI_API_KEY`.\n",
    "\n",
    "Below, we validate that this key is set properly in your environment.\n",
    "\n",
    "**Packages**\n",
    "\n",
    "To interface with both UnityCatalog and AutoGen, you will need to install the following package:\n",
    "\n",
    "```shell\n",
    "pip install unitycatalog-autogen\n",
    "```\n",
    "\n",
    "**Note**\n",
    "The official Microsoft AutoGen package has been renamed from `pyautogen` to `autogen-agentchat`. \n",
    "There are additional forked version of the AutoGen package that are not contributed by Microsoft and will not work with this integration. \n",
    "For further information, please see the [official clarification statement](https://github.com/microsoft/autogen/discussions/4217). \n",
    "\n",
    "**Note**\n",
    "The example shown here is for `autogen` 0.4.0 and above. Earlier versions of the library will not function with the Unity Catalog AI Toolkit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "72b76dbc-f999-4487-a405-a44ae38da028",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "assert \"OPENAI_API_KEY\" in os.environ, (\n",
    "    \"Please set the OPENAI_API_KEY environment variable to your OpenAI API key\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc8f4354-33ff-4475-b97d-6024bdd53044",
   "metadata": {},
   "source": [
    "## Configuration and Client setup\n",
    "\n",
    "In order to connect to your Unity Catalog server, you'll need an instance of the `ApiClient` from the `unitycatalog-client` package. \n",
    "\n",
    "> Note: If you don't already have a Catalog and a Schema created, be sure to create them before running this notebook and adjust the `CATALOG` and `SCHEMA` variables below to suit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "337293cd-0028-47a8-8e4b-e14fcea2c28f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from unitycatalog.ai.autogen.toolkit import UCFunctionToolkit\n",
    "from unitycatalog.ai.core.client import UnitycatalogFunctionClient\n",
    "from unitycatalog.client import ApiClient, Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "85f843c0-2691-4f88-8eb3-322e1b8c7a3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = Configuration()\n",
    "config.host = \"http://localhost:8080/api/2.1/unity-catalog\"\n",
    "\n",
    "# The base ApiClient is async\n",
    "api_client = ApiClient(configuration=config)\n",
    "\n",
    "client = UnitycatalogFunctionClient(api_client=api_client)\n",
    "\n",
    "CATALOG = \"AICatalog\"\n",
    "SCHEMA = \"AISchema\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "278fbac1-34de-487e-8535-5595d77222ad",
   "metadata": {},
   "source": [
    "## Define a function and register it to Unity Catalog\n",
    "\n",
    "In this next section, we'll be defining some Python functions and creating them within Unity Catalog so that they can be retrieved and used as tools by an AutoGen agent. \n",
    "\n",
    "There are a few things to keep in mind when creating functions for use with the `create_python_function` API:\n",
    "\n",
    "- Ensure that your have properly defined types for all arguments and for the return of the function.\n",
    "- Ensure that you have a Google-style docstring defined that includes descriptions for the function, each argument, and the return of the function. This is critical, as these are used to populate the metadata associated with the function within Unity Catalog, providing contextual data for an LLM to understand when and how to call the tool associated with this function.\n",
    "- If there are packages being called that are not part of core Python, ensure that the import statements are locally scoped (defined within the function body)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cfc1dd58-c1a5-4d64-b1d1-b277e4516f30",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_vpd(temperature_c: float, dew_point_c: float) -> float:\n",
    "    \"\"\"\n",
    "    Calculate Vapor Pressure Deficit (VPD) given temperature and dew point in Celsius.\n",
    "\n",
    "    Args:\n",
    "        temperature_c (float): Air temperature in Celsius.\n",
    "        dew_point_c (float): Dew point temperature in Celsius.\n",
    "\n",
    "    Returns:\n",
    "        Vapor Pressure Deficit in hPa.\n",
    "    \"\"\"\n",
    "    import math  # local imports are needed for execution\n",
    "\n",
    "    a = 17.625\n",
    "    b = 243.04\n",
    "\n",
    "    e_s = 6.1094 * math.exp((a * temperature_c) / (b + temperature_c))\n",
    "    e_a = 6.1094 * math.exp((a * dew_point_c) / (b + dew_point_c))\n",
    "\n",
    "    vpd = e_s - e_a\n",
    "    return vpd\n",
    "\n",
    "\n",
    "def fahrenheit_to_celsius(fahrenheit: float) -> float:\n",
    "    \"\"\"\n",
    "    Converts temperature from Fahrenheit to Celsius.\n",
    "\n",
    "    Args:\n",
    "        fahrenheit (float): Temperature in degrees Fahrenheit.\n",
    "\n",
    "    Returns:\n",
    "        float: Temperature in degrees Celsius.\n",
    "    \"\"\"\n",
    "    return (fahrenheit - 32) * 5.0 / 9.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "92c800f2-7891-42ad-8bd8-2b3da7b40a32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionInfo(name='calculate_vpd', catalog_name='AICatalog', schema_name='AISchema', input_params=FunctionParameterInfos(parameters=[FunctionParameterInfo(name='temperature_c', type_text='DOUBLE', type_json='{\"name\": \"temperature_c\", \"type\": \"double\", \"nullable\": false, \"metadata\": {\"comment\": \"Air temperature in Celsius.\"}}', type_name=<ColumnTypeName.DOUBLE: 'DOUBLE'>, type_precision=None, type_scale=None, type_interval_type=None, position=0, parameter_mode=None, parameter_type=None, parameter_default=None, comment='Air temperature in Celsius.'), FunctionParameterInfo(name='dew_point_c', type_text='DOUBLE', type_json='{\"name\": \"dew_point_c\", \"type\": \"double\", \"nullable\": false, \"metadata\": {\"comment\": \"Dew point temperature in Celsius.\"}}', type_name=<ColumnTypeName.DOUBLE: 'DOUBLE'>, type_precision=None, type_scale=None, type_interval_type=None, position=1, parameter_mode=None, parameter_type=None, parameter_default=None, comment='Dew point temperature in Celsius.')]), data_type=<ColumnTypeName.DOUBLE: 'DOUBLE'>, full_data_type='DOUBLE', return_params=None, routine_body='EXTERNAL', routine_definition='import math  # local imports are needed for execution\\n\\na = 17.625\\nb = 243.04\\n\\ne_s = 6.1094 * math.exp((a * temperature_c) / (b + temperature_c))\\ne_a = 6.1094 * math.exp((a * dew_point_c) / (b + dew_point_c))\\n\\nvpd = e_s - e_a\\nreturn vpd', routine_dependencies=None, parameter_style='S', is_deterministic=True, sql_data_access='NO_SQL', is_null_call=False, security_type='DEFINER', specific_name='calculate_vpd', comment='Calculate Vapor Pressure Deficit (VPD) given temperature and dew point in Celsius.', properties='null', full_name='AICatalog.AISchema.calculate_vpd', owner=None, created_at=1737060978247, created_by=None, updated_at=1737060978247, updated_by=None, function_id='2813240c-9a58-4fe2-9276-71a746064eee', external_language='PYTHON')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.create_python_function(func=calculate_vpd, catalog=CATALOG, schema=SCHEMA, replace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "aa0ed262-ddac-461f-acaf-0265711498b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FunctionInfo(name='fahrenheit_to_celsius', catalog_name='AICatalog', schema_name='AISchema', input_params=FunctionParameterInfos(parameters=[FunctionParameterInfo(name='fahrenheit', type_text='DOUBLE', type_json='{\"name\": \"fahrenheit\", \"type\": \"double\", \"nullable\": false, \"metadata\": {\"comment\": \"Temperature in degrees Fahrenheit.\"}}', type_name=<ColumnTypeName.DOUBLE: 'DOUBLE'>, type_precision=None, type_scale=None, type_interval_type=None, position=0, parameter_mode=None, parameter_type=None, parameter_default=None, comment='Temperature in degrees Fahrenheit.')]), data_type=<ColumnTypeName.DOUBLE: 'DOUBLE'>, full_data_type='DOUBLE', return_params=None, routine_body='EXTERNAL', routine_definition='return (fahrenheit - 32) * 5.0 / 9.0', routine_dependencies=None, parameter_style='S', is_deterministic=True, sql_data_access='NO_SQL', is_null_call=False, security_type='DEFINER', specific_name='fahrenheit_to_celsius', comment='Converts temperature from Fahrenheit to Celsius.', properties='null', full_name='AICatalog.AISchema.fahrenheit_to_celsius', owner=None, created_at=1737060978267, created_by=None, updated_at=1737060978267, updated_by=None, function_id='285a59ec-721d-4746-bb08-8fa2fc745339', external_language='PYTHON')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.create_python_function(\n",
    "    func=fahrenheit_to_celsius, catalog=CATALOG, schema=SCHEMA, replace=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e28d0c6d-7af4-4daf-8993-d5037524c2de",
   "metadata": {},
   "source": [
    "## Create a Toolkit instance of the function(s)\n",
    "\n",
    "Now that the function has been created within Unity Catalog, we can use the `unitycatalog-autogen` package to create a toolkit instance that our Agent will 'understand' as a valid tool in its APIs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0399a04e-b5da-4e7c-a406-71dc99f116ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "vpd_tool = f\"{CATALOG}.{SCHEMA}.calculate_vpd\"\n",
    "f_to_c_tool = f\"{CATALOG}.{SCHEMA}.fahrenheit_to_celsius\"\n",
    "\n",
    "toolkit = UCFunctionToolkit(function_names=[vpd_tool, f_to_c_tool], client=client)\n",
    "\n",
    "tools = toolkit.tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a98eff1a-50d5-445e-bac2-57a3c0c4e7b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<unitycatalog.ai.autogen.toolkit.UCFunctionTool at 0x10c1b4790>,\n",
       " <unitycatalog.ai.autogen.toolkit.UCFunctionTool at 0x109c74990>]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64d0b9b4-b8d6-461e-9c50-fb852cbd670d",
   "metadata": {},
   "source": [
    "## Create a Conversable Agent that uses our tool\n",
    "\n",
    "Now we get to actually create an Agent. As part of our definition, we'll be applying the tool defintion from our Toolkit instance. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "af5bcc69-809a-4e83-843d-4102741fd1ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from autogen_agentchat.agents import AssistantAgent\n",
    "from autogen_agentchat.messages import TextMessage\n",
    "from autogen_core import CancellationToken\n",
    "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
    "\n",
    "OPENAI_API_KEY = os.environ[\"OPENAI_API_KEY\"]\n",
    "\n",
    "model_client = OpenAIChatCompletionClient(\n",
    "    model=\"gpt-4o\",\n",
    "    api_key=OPENAI_API_KEY,\n",
    "    seed=222,\n",
    "    temperature=0,\n",
    ")\n",
    "\n",
    "weather_agent = AssistantAgent(\n",
    "    name=\"assistant\",\n",
    "    system_message=\"You are a helpful AI assistant that specializes in answering questions about weather phenomena. \"\n",
    "    \"If there are tools available to perform these calculations, please use them and ensure that you are operating \"\n",
    "    \"with validated data calculations.\",\n",
    "    model_client=model_client,\n",
    "    tools=tools,\n",
    "    reflect_on_tool_use=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49ebacc1-dea5-4252-b8bf-ffb131a6647f",
   "metadata": {},
   "source": [
    "## Ask the Agent a question\n",
    "\n",
    "Now that we have everything configured, let's test out our Agent! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e459e12d-1d36-46c7-b96a-5418fa8e7a72",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TextMessage(source='assistant', models_usage=RequestUsage(prompt_tokens=286, completion_tokens=14), content='973.2°F is approximately 522.89°C.', type='TextMessage')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_input = \"I need some help converting 973.2F to Celsius.\"\n",
    "\n",
    "response = await weather_agent.on_messages(\n",
    "    [TextMessage(content=user_input, source=\"User\")], CancellationToken()\n",
    ")\n",
    "\n",
    "response.chat_message"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mlflow-311",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
